"""
期货数据采集器模块

负责通过akshare API获取期货数据，进行清洗，并保存到数据库
"""
import time
import datetime
import logging
import akshare as ak
import pandas as pd
from sqlalchemy.orm import Session
from sqlalchemy.exc import IntegrityError

from modules.models.base import get_db
from modules.models.stock import SecurityInfo, DailyQuote, CollectionTask
from config.settings import DATA_COLLECTION

# 配置日志
logger = logging.getLogger("futures_collector")


class FuturesCollector:
    """期货数据采集类"""
    
    def __init__(self):
        """初始化采集器"""
        self.retry_limit = DATA_COLLECTION["retry_limit"]
        self.retry_delay = DATA_COLLECTION["retry_delay"]
        self.batch_size = DATA_COLLECTION["batch_size"]
        self.db = next(get_db())
    
    def __del__(self):
        """析构函数，确保关闭数据库连接"""
        if hasattr(self, 'db') and self.db:
            self.db.close()
    
    def get_futures_exchanges(self):
        """获取期货交易所列表"""
        # 主要期货交易所列表
        exchanges = [
            {"code": "SHFE", "name": "上海期货交易所", "api_func": "futures_shfe_daily_sina"},
            {"code": "DCE", "name": "大连商品交易所", "api_func": "futures_dce_daily_sina"},
            {"code": "CZCE", "name": "郑州商品交易所", "api_func": "futures_czce_daily_sina"},
            {"code": "CFFEX", "name": "中国金融期货交易所", "api_func": "futures_cffex_daily_sina"},
            {"code": "INE", "name": "上海国际能源交易中心", "api_func": "futures_ine_daily_sina"},
            {"code": "GFEX", "name": "广州期货交易所", "api_func": "futures_gfex_daily_sina"}
        ]
        return exchanges
    
    def get_futures_symbols(self, exchange_code=None):
        """
        获取期货合约代码列表
        
        Args:
            exchange_code: 交易所代码，为None时返回所有交易所
        
        Returns:
            合约代码列表
        """
        try:
            logger.info(f"开始获取期货合约代码，交易所：{exchange_code or '所有'}")
            
            # 根据交易所获取合约列表
            if exchange_code == "SHFE":
                # 上海期货交易所
                df = ak.futures_shfe_realtime()
                df = df.rename(columns={"symbol": "code", "name": "name"})
                df["exchange"] = "SHFE"
            elif exchange_code == "DCE":
                # 大连商品交易所
                df = ak.futures_dce_position_rank()
                df = df[["variety", "code"]].drop_duplicates()
                df = df.rename(columns={"variety": "name"})
                df["exchange"] = "DCE"
            elif exchange_code == "CZCE":
                # 郑州商品交易所
                df = ak.futures_czce_warehouse_receipt()
                df = df[["var", "code"]].drop_duplicates()
                df = df.rename(columns={"var": "name"})
                df["exchange"] = "CZCE"
            elif exchange_code == "CFFEX":
                # 中国金融期货交易所
                df = ak.futures_contract_detail(symbol="IF")
                df = df[["code", "instrument_name"]].drop_duplicates()
                df = df.rename(columns={"instrument_name": "name"})
                df["exchange"] = "CFFEX"
            else:
                # 获取所有交易所的期货主力合约
                df = ak.futures_main_sina()
                df = df[["symbol", "name", "exchange"]].drop_duplicates()
                df = df.rename(columns={"symbol": "code"})
            
            logger.info(f"成功获取{len(df)}个期货合约代码")
            return df
        
        except Exception as e:
            logger.error(f"获取期货合约代码失败: {str(e)}")
            raise
    
    def collect_futures_info(self, exchange_code=None):
        """
        采集并更新期货合约基本信息
        
        Args:
            exchange_code: 交易所代码，为None时采集所有交易所
        
        Returns:
            (新增数量, 更新数量)
        """
        logger.info(f"开始采集期货合约基本信息，交易所：{exchange_code or '所有'}")
        
        try:
            # 获取期货合约列表
            futures_df = self.get_futures_symbols(exchange_code)
            
            # 获取已存在的期货代码
            existing_codes = {
                code[0] for code in 
                self.db.query(SecurityInfo.code).filter(SecurityInfo.type == 'future').all()
            }
            
            # 准备批量插入或更新的数据
            new_futures = []
            update_futures = []
            
            for _, row in futures_df.iterrows():
                code = row['code']
                name = row['name']
                exchange = row.get('exchange', '')
                
                # 检查是否存在
                if code in existing_codes:
                    # 更新现有合约
                    future = self.db.query(SecurityInfo).filter(
                        SecurityInfo.code == code,
                        SecurityInfo.type == 'future'
                    ).first()
                    
                    if future:
                        future.name = name
                        future.exchange = exchange
                        future.update_time = datetime.datetime.now()
                        update_futures.append(future)
                else:
                    # 创建新合约记录
                    new_future = SecurityInfo(
                        code=code,
                        name=name,
                        type='future',
                        exchange=exchange,
                        is_active=1
                    )
                    new_futures.append(new_future)
            
            # 批量保存新合约
            if new_futures:
                self.db.add_all(new_futures)
                self.db.commit()
                logger.info(f"成功添加{len(new_futures)}个新期货合约")
            
            # 批量更新现有合约
            if update_futures:
                self.db.commit()
                logger.info(f"成功更新{len(update_futures)}个期货合约信息")
            
            return len(new_futures), len(update_futures)
        
        except Exception as e:
            self.db.rollback()
            logger.error(f"采集期货合约基本信息失败: {str(e)}")
            raise
    
    def collect_futures_daily_data(self, symbol, start_date=None, end_date=None, task_id=None):
        """
        采集指定期货合约的日线数据
        
        Args:
            symbol: 期货合约代码
            start_date: 开始日期，默认为None
            end_date: 结束日期，默认为None
            task_id: 任务ID
        
        Returns:
            成功保存的数据条数
        """
        # 更新任务状态
        if task_id:
            self._update_task_status(task_id, "running")
        
        # 获取合约信息
        security = self.db.query(SecurityInfo).filter(
            SecurityInfo.code == symbol,
            SecurityInfo.type == 'future'
        ).first()
        
        if not security:
            error_msg = f"期货合约 {symbol} 不存在"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
        
        # 设置日期范围
        if not start_date:
            # 从最后一条记录的下一天开始
            last_record = self.db.query(DailyQuote).filter(
                DailyQuote.security_id == security.id
            ).order_by(DailyQuote.trade_date.desc()).first()
            
            if last_record:
                start_date = (last_record.trade_date + datetime.timedelta(days=1)).strftime('%Y%m%d')
            else:
                # 默认采集近1年数据
                start_date = (datetime.datetime.now() - datetime.timedelta(days=365)).strftime('%Y%m%d')
        
        if not end_date:
            end_date = datetime.datetime.now().strftime('%Y%m%d')
        
        logger.info(f"开始采集期货合约 {symbol} 的日线数据，时间范围: {start_date} - {end_date}")
        
        try:
            # 根据交易所选择不同的接口
            if security.exchange == "SHFE":
                method = "futures_shfe_daily_sina"
            elif security.exchange == "DCE":
                method = "futures_dce_daily_sina"
            elif security.exchange == "CZCE":
                method = "futures_czce_daily_sina"
            elif security.exchange == "CFFEX":
                method = "futures_cffex_daily_sina"
            elif security.exchange == "INE":
                method = "futures_ine_daily_sina"
            elif security.exchange == "GFEX":
                method = "futures_gfex_daily_sina"
            else:
                # 默认使用新浪接口
                method = "futures_daily_sina"
            
            # 重试机制
            retry_count = 0
            df = None
            
            while retry_count < self.retry_limit:
                try:
                    # 调用akshare接口
                    if method == "futures_daily_sina":
                        df = ak.futures_daily_sina(symbol=symbol)
                    else:
                        # 动态调用对应的交易所接口
                        api_func = getattr(ak, method)
                        df = api_func()
                        # 筛选指定合约的数据
                        if 'symbol' in df.columns:
                            df = df[df['symbol'] == symbol]
                        elif 'code' in df.columns:
                            df = df[df['code'] == symbol]
                        elif 'var_name' in df.columns:
                            df = df[df['var_name'] == symbol]
                    break
                except Exception as e:
                    retry_count += 1
                    logger.warning(f"采集期货合约 {symbol} 数据失败，重试 {retry_count}/{self.retry_limit}: {str(e)}")
                    if retry_count >= self.retry_limit:
                        raise
                    time.sleep(self.retry_delay)
            
            if df is None or df.empty:
                logger.info(f"期货合约 {symbol} 在指定时间范围内没有数据")
                if task_id:
                    self._update_task_status(task_id, "success", None, 0)
                return 0
            
            # 转换和清洗数据
            df = self._clean_futures_data(df)
            
            # 筛选日期范围
            if 'trade_date' in df.columns:
                df['trade_date'] = pd.to_datetime(df['trade_date'])
                start_date_dt = pd.to_datetime(start_date)
                end_date_dt = pd.to_datetime(end_date)
                df = df[(df['trade_date'] >= start_date_dt) & (df['trade_date'] <= end_date_dt)]
            
            # 保存到数据库
            inserted_count = self._save_futures_quotes(security.id, df)
            
            # 更新任务状态
            if task_id:
                self._update_task_status(task_id, "success", None, inserted_count)
            
            return inserted_count
            
        except Exception as e:
            error_msg = f"采集期货合约 {symbol} 数据失败: {str(e)}"
            logger.error(error_msg)
            if task_id:
                self._update_task_status(task_id, "failed", error_msg)
            return 0
    
    def _clean_futures_data(self, df):
        """清洗和标准化期货数据"""
        if df.empty:
            return df
        
        # 标准化列名
        df.columns = [col.lower() for col in df.columns]
        
        # 期货数据的列名映射
        name_mapping = {
            'date': 'trade_date',
            'open_price': 'open',
            'highest_price': 'high',
            'lowest_price': 'low',
            'close_price': 'close',
            'volume': 'volume',
            'position': 'open_interest',
            'turnover': 'amount'
        }
        
        # 应用列名映射
        df = df.rename(columns={k: v for k, v in name_mapping.items() if k in df.columns})
        
        # 处理日期
        if 'trade_date' in df.columns:
            df['trade_date'] = pd.to_datetime(df['trade_date'])
        
        # 处理数值类型
        numeric_cols = ['open', 'high', 'low', 'close', 'volume', 'amount', 'open_interest']
        for col in numeric_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
        
        # 保留需要的列
        keep_cols = [col for col in ['trade_date', 'open', 'high', 'low', 'close', 
                                      'volume', 'amount', 'open_interest'] if col in df.columns]
        
        return df[keep_cols]
    
    def _save_futures_quotes(self, security_id, df):
        """保存期货日线数据到数据库"""
        if df.empty:
            return 0
        
        inserted_count = 0
        
        try:
            # 获取已存在的交易日期
            existing_dates = {
                date[0] for date in 
                self.db.query(DailyQuote.trade_date).filter(
                    DailyQuote.security_id == security_id
                ).all()
            }
            
            # 准备批量插入
            batch_records = []
            
            for _, row in df.iterrows():
                trade_date = row['trade_date'].date()
                
                # 跳过已存在的记录
                if trade_date in existing_dates:
                    continue
                
                daily_quote = DailyQuote(
                    security_id=security_id,
                    trade_date=trade_date,
                    open=row.get('open'),
                    high=row.get('high'),
                    low=row.get('low'),
                    close=row.get('close'),
                    adj_close=row.get('close'),  # 期货没有复权价
                    volume=row.get('volume'),
                    amount=row.get('amount'),
                    turnover=None  # 期货没有换手率
                )
                
                batch_records.append(daily_quote)
                inserted_count += 1
                
                # 分批提交
                if len(batch_records) >= self.batch_size:
                    self.db.add_all(batch_records)
                    self.db.commit()
                    batch_records = []
            
            # 提交剩余记录
            if batch_records:
                self.db.add_all(batch_records)
                self.db.commit()
            
            logger.info(f"成功保存 {inserted_count} 条期货日线数据")
            return inserted_count
            
        except IntegrityError as e:
            self.db.rollback()
            logger.error(f"保存期货日线数据失败(完整性错误): {str(e)}")
            return 0
        except Exception as e:
            self.db.rollback()
            logger.error(f"保存期货日线数据失败: {str(e)}")
            raise
    
    def _update_task_status(self, task_id, status, error_message=None, result_count=None):
        """更新任务状态"""
        try:
            task = self.db.query(CollectionTask).filter(CollectionTask.task_id == task_id).first()
            if task:
                task.status = status
                if error_message:
                    task.error_message = error_message
                if result_count is not None:
                    task.result_count = result_count
                task.update_time = datetime.datetime.now()
                self.db.commit()
        except Exception as e:
            self.db.rollback()
            logger.error(f"更新任务状态失败: {str(e)}")


# 测试函数
def test_futures_collector():
    """测试期货数据采集器功能"""
    collector = FuturesCollector()
    
    # 1. 测试获取期货交易所列表
    exchanges = collector.get_futures_exchanges()
    print(f"成功获取 {len(exchanges)} 个期货交易所")
    for exchange in exchanges:
        print(f"交易所: {exchange['name']} ({exchange['code']})")
    
    # 2. 测试获取期货合约列表
    try:
        futures_list = collector.get_futures_symbols()
        print(f"成功获取 {len(futures_list)} 个期货合约")
        print(futures_list.head())
    except Exception as e:
        print(f"获取期货合约列表失败: {str(e)}")
    
    # 3. 测试采集期货合约基本信息
    try:
        new_count, update_count = collector.collect_futures_info()
        print(f"新增期货合约: {new_count}, 更新期货合约: {update_count}")
    except Exception as e:
        print(f"采集期货合约基本信息失败: {str(e)}")
    
    # 4. 测试采集日线数据
    if futures_list is not None and not futures_list.empty:
        # 选择第一个期货合约进行测试
        test_symbol = futures_list.iloc[0]['code']
        try:
            # 采集最近30天数据
            end_date = datetime.datetime.now().strftime('%Y%m%d')
            start_date = (datetime.datetime.now() - datetime.timedelta(days=30)).strftime('%Y%m%d')
            
            count = collector.collect_futures_daily_data(test_symbol, start_date, end_date)
            print(f"期货合约 {test_symbol} 成功采集 {count} 条日线数据")
        except Exception as e:
            print(f"采集期货合约 {test_symbol} 日线数据失败: {str(e)}")


if __name__ == "__main__":
    # 配置日志
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 执行测试
    test_futures_collector()